from threading import Thread
from utils.common import exec_shell
import time
from influxdb_client import Point
import logging
import os
from utils.common import init_log
import subprocess

def collect_cpu_utility_or_memory_usage(period,cmd,write_api,name,fn,pid):
    while True:
        output,_=exec_shell(cmd)
        res=output.decode().strip()
        write_api.write(bucket='btresults',record=[Point(name).tag('type','host').field(fn,res)])
        logging.info(f'{fn}: {res}')
        time.sleep(period)
        try:
            os.kill(pid,0)
        except:
            break

def collect_cpu_utility_and_memory_usage(period,cmd,write_api,name,pid):
    while True:
        output,_=exec_shell(cmd)
        cpu,mem=output.decode().strip().split()
        write_api.write(bucket='btresults',record=[Point(name).tag('type','host').field('cpu_utility',cpu).field('memory_usage',mem)])
        logging.info(f'cpu_utility: {cpu},memory_usage: {mem}')
        time.sleep(period)
        try:
            os.kill(pid,0)
        except:
            break

def collect_gpu_info(period,cmd,write_api,name,fns,pid):
    while True:
        # print(cmd)
        output,_=exec_shell(cmd)
        rs=output.decode().strip().split(', ')
        if len(rs)>0:
            _point=Point(name).tag('type','nvidia_gpu')
            for i in range(len(rs)):
                _point=_point.field(fns[i],rs[i])
            write_api.write(bucket='btresults',record=[_point])
            logging.info(f'{fns}: {rs}')
        time.sleep(period)
        try:
            os.kill(pid,0)
        except:
            break
        
def collect_gpu_accounted_info(period,cmd,write_api,name,fns,pid):
    while True:
        # print('--------collect-gpu=accounted-info',cmd)
        output,err=exec_shell(cmd)
        # print(output,err)
        if output!=-1000:
            rs=output.decode().strip().split(', ')
            if len(rs)>1:
                _point=Point(name).tag('type','nvidia_gpu_accounted')
                for i in range(1,len(rs)):
                    _point=_point.field(fns[i],rs[i])
                write_api.write(bucket='btresults',record=[_point])
                logging.info(f'{fns}: {rs}')
        time.sleep(period)
        try:
            os.kill(pid,0)
        except:
            break

def monitor_pid(pid):
    while True:
        try:
            os.kill(pid,0)
        except:
            break
        time.sleep(1)

def collect_hardware_cost_metric_per_period(conn,bttask,pid,write_api,name,metrics):
    init_log(f'logs/{name}.log')
    threads=[]
    if all(x in metrics for x in ['cpu_utility','memory_usage']):
        threads.append(Thread(target=collect_cpu_utility_and_memory_usage,args=(bttask.hardware_cost_collection_interval,f"top -bn 1 -p {pid}|tail -3|tail -1|awk '{{cpu=NF-3}} {{mem=NF-6}} {{print $cpu\" \"$mem}}'",write_api,name,pid)))
    elif any(x in metrics for x in ['cpu_utility','memory_usage']):
        if 'memory_usage' in metrics:
            threads.append(Thread(target=collect_cpu_utility_or_memory_usage,args=(bttask.hardware_cost_collection_interval,f"top -bn 1 -p {pid}|tail -3|tail -1|awk '{{mem=NF-6}} {{print $mem}}'",write_api,name,'memory_usage',pid)))
        else:
            threads.append(Thread(target=collect_cpu_utility_or_memory_usage,args=(bttask.hardware_cost_collection_interval,f"top -bn 1 -p {pid}|tail -3|tail -1|awk '{{cpu=NF-3}} {{print $cpu}}'",write_api,name,'cpu_utility',pid)))
    query_gpu_str=''
    query_accounted_apps_str=''
    fns=[]
    accounted_fns=[pid]
    if 'gpu_utility' in metrics:
        query_gpu_str+='utilization.gpu,'
        fns.append('gpu_utility')
        query_accounted_apps_str+='gpu_util,'
        accounted_fns.append('gpu_util')
    if 'gpu_memory_utility' in metrics:
        query_gpu_str+='utilization.memory,'
        fns.append('gpu_memory_utility')
        query_accounted_apps_str+='mem_util,'
        accounted_fns.append('mem_util')
    if 'gpu_memory_usage' in metrics:
        query_accounted_apps_str+='max_memory_usage,'
        accounted_fns.append('gpu_memory_usage')
    if 'gpu_temperature' in metrics:
        query_gpu_str+='temperature.gpu,'
        fns.append('gpu_temperature')
    if 'gpu_power' in metrics:
        query_gpu_str+='power.draw,'
        fns.append('gpu_power')
    if 'gpu_clock_frequency' in metrics:
        query_gpu_str+='clocks.sm,'
        fns.append('gpu_clock_frequency')
    if query_gpu_str!='':
        threads.append(Thread(target=collect_gpu_info,args=(bttask.hardware_cost_collection_interval,f'nvidia-smi --query-gpu={query_gpu_str[:-1]} --format=csv,noheader',write_api,name,fns,pid)))
    if query_accounted_apps_str!='':
        exec_shell('nvidia-smi -caa')
        threads.append(Thread(target=collect_gpu_accounted_info,args=(bttask.hardware_cost_collection_interval,f"nvidia-smi --query-accounted-apps=pid,{query_accounted_apps_str[:-1]} --format=csv,noheader|grep '{pid}, '",write_api,name,fns,pid)))
    conn.send('ready')
    conn.recv()
    for t in threads:
        t.daemon=True
        t.start()
    monitor_pid(pid)